"""
Author: Morphlng
Date: 2024-04-08 13:19:15
LastEditTime: 2024-12-16 23:08:58
LastEditors: Morphlng
Description: Handle the retrieval and processing of observation for an agent
FilePath: /DrivingGym/src/driving_gym/environment/agent/obs/obs_handler.py
"""

from __future__ import annotations

import os
from collections import defaultdict
from typing import TYPE_CHECKING

import gymnasium as gym
import numpy as np

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.obs.base_obs import BaseObs
from driving_gym.environment.scenario.base_scenario import BaseScenario
from driving_gym.misc.util import dynamic_import

if TYPE_CHECKING:
    from driving_gym.simulation.adapter_interface import AdapterInterface


class ObsHandler:
    """
    Example of config:

    ```python
    obs_config = {
        "framestack": 1,
        "frameskip": 4,
        "items": {
            "image": {
                "type": "image_obs",
                "source": "hero:rgb", # or "rgb"
                "width": 256,
                "height": 256,
            },
            "semantic": {
                "type": "kinematic_obs",
                "actor_count": 5,
            }
        }
    }
    ```
    """

    def __init__(self, config: dict, adapter: "AdapterInterface"):
        self.config = config
        self.frame_stack = config.get("framestack", 1)
        self.frame_skip = config.get("frameskip", 0)
        self.handlers = {
            name: self._create_handler(conf, adapter)
            for name, conf in config["items"].items()
        }

    def reset(self):
        for handler in self.handlers.values():
            handler.reset()

    def get_obs(self, snapshots: list[Snapshot], scenario: BaseScenario):
        if len(snapshots) < self.frame_stack:
            raise ValueError("Not enough snapshots for frame stack")

        obs_per_frame = []
        for snapshot in snapshots:
            obs_per_frame.append(
                {
                    name: handler.get_obs(snapshot, scenario)
                    for name, handler in self.handlers.items()
                }
            )
        stack_obs = self._stack_obs(obs_per_frame)
        return stack_obs

    def get_observation_space(self):
        obs_space = {
            name: handler.get_observation_space()
            for name, handler in self.handlers.items()
        }

        if self.frame_stack > 1:
            for name, space in obs_space.items():
                dtype = space.dtype
                shape = (self.frame_stack, *space.shape)
                low = np.full(shape, space.low, dtype=dtype)
                high = np.full(shape, space.high, dtype=dtype)
                obs_space[name] = gym.spaces.Box(low, high, shape, dtype)

        return gym.spaces.Dict(obs_space)

    def _create_handler(self, config: dict, adapter: "AdapterInterface") -> BaseObs:
        obs_type = config["type"]
        root = os.path.join(os.path.dirname(__file__))
        handler = dynamic_import(obs_type, root, args=(config, adapter), suffix="_obs")
        return handler

    def _stack_obs(self, obs_per_frame: list[dict[str, np.ndarray]]):
        if len(obs_per_frame) == 1:
            return obs_per_frame[0]

        final_obs = defaultdict(list)
        for obs in obs_per_frame:
            for name, value in obs.items():
                final_obs[name].append(value)
        return {name: np.stack(final_obs[name], axis=0) for name in final_obs}
